// Copyright (C) 2023 Computer Vision Lab, Electrical Engineering, 
// Indian Institute of Science, Bengaluru, India.
// All rights reserved.
// 
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//     * Redistributions of source code must retain the above copyright
//       notice, this list of conditions and the following disclaimer.
//     * Redistributions in binary form must reproduce the above
//       copyright notice, this list of conditions and the following
//       disclaimer in the documentation and/or other materials provided
//       with the distribution.
//     * Neither the name of Indian Institute of Science nor the
//       names of its contributors may be used to endorse or promote products
//       derived from this software without specific prior written permission.
// 
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS 
// OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL 
// THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 
// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER 
// DEALINGS IN THE SOFTWARE.
// 
// 
// This file is a part of the implementation for the paper:
// Chitturi Sidhatha, Lalit Manam and Venu Madhav Govindu, 
// Adaptive Annealing for Robust Geometric Estimation,
// Proceedings of the Computer Vision and Pattern Recognition 2023
// 
// This code base is adapted from the code release of Fast Global Registration,
// Qian-Yi Zhou, Jaesik Park, and Vladlen Koltun, ECCV 2016

#include "app.h"
#include <omp.h>

using namespace Eigen;
using namespace std;
using namespace gom;

// New functions added at the beginning only

Eigen::Matrix4d CApp::WeightedUmeyama(const Eigen::MatrixXd& ptsX,const Eigen::MatrixXd& ptsY,const Eigen::VectorXd& weights)
{
    /* Weighted Umeyama */
    // Let ptsX, ptsY be the given point matrices such that ptsY = s*R*ptsX+T, 

    if(ptsX.rows()!=ptsY.rows() || ptsX.cols()!=ptsY.cols())
    {
        std::cout<< "Dimensions of matched correspondences are not the same. Expected dimension is mxN\n";
		exit(1);
    }

    bool compute_scale = false; //scale is not computed.
    int N = ptsX.cols();
    int m = ptsX.rows();

    double cn = weights.sum();

    Eigen::VectorXd mu_x = ((ptsX.array().rowwise())*(weights.transpose().array())).rowwise().sum()/cn;
    Eigen::VectorXd mu_y = ((ptsY.array().rowwise())*(weights.transpose().array())).rowwise().sum()/cn;
    
    double sig_x2 = weights.dot(((((ptsX.colwise()-mu_x).array().square()).colwise().sum()).matrix()))/cn;
    double sig_y2 = weights.dot(((((ptsY.colwise()-mu_y).array().square()).colwise().sum()).matrix()))/cn;
    
    Eigen::Matrix3d sig_xy = ((ptsY.array().rowwise())*(weights.transpose().array())).matrix()*(ptsX.transpose())/cn- mu_y*(mu_x.transpose());

    //Computing SVD
    Eigen::JacobiSVD<Eigen::Matrix3d, Eigen::ComputeFullU | Eigen::ComputeFullV> svd(sig_xy);
	Eigen::VectorXd S = Eigen::VectorXd::Ones(m);
    if(((svd.matrixU()*svd.matrixV()).determinant())<0)
        S(m-1) = -1;// S(2,2)= -1;

	Eigen::MatrixXd M = Eigen::MatrixXd::Identity(m+1,m+1);    
    M.block(0,0,m,m).noalias() = svd.matrixU()*S.asDiagonal()*(svd.matrixV().transpose());
    double scale;
    if(compute_scale)
    {
		scale = (svd.singularValues().dot(S))/sig_x2;
    }
    else
    {
        scale = 1.0d; // For now, in the output we are not giving scale.
    }
    M.col(m).head(m) = mu_y-scale*M.topLeftCorner(m,m)*mu_x;
	return M;
}


void CApp::GetTensorValue()
{
	Eigen::Tensor<double, 4> t(1, 2, 3, 4);
	t(0, 1, 2, 3) = 42.0;
	std::cout << t << std::endl;
}
void CApp::LoadPointCloud(const PointsMat& pts1, const PointsMat& pts2)
{
	// pointcloud_.push_back(pts);
	ptsA_ = pts1.transpose();
	ptsB_ = pts2.transpose();
}

void CApp::SetMatches()
{
	int fi = 0;
	// int fj = 1;
	std::cout << "The number of matches are:" << ptsA_.cols() << "of dimension" << ptsA_.rows() << std::endl;
	int nPti = ptsA_.cols();
	// int nPtj = pointcloud_[fj].size();  // Since we assumed that the input has only matched correspondences, so they will be of equal size.
	for (int i=0; i < nPti; i++)
		corres_.push_back(std::pair<int,int>(i,i));
	num_corr_ = nPti;
}

// template <typename T>
static void gom::sort_indexes(const Eigen::VectorXd &v, Eigen::VectorXi& idx) {
	// initialize original index locations
	iota(idx.begin(), idx.end(), 0);
	stable_sort(idx.begin(), idx.end(), [&v](int i1, int i2) {return v(i1) < v(i2);});
}	


/* Unequal probability sampling; with-replacement case 
 * n are the lengths of p and perm. p contains probabilities, perm
 * contains the actual outcomes, and ans contains an array of values 
 * that were sampled.
 * Source: https://stats.stackexchange.com/questions/67911/how-to-sample-from-a-discrete-distribution
 */

static void gom::ProbSampleReplace(Eigen::VectorXd& p, Eigen::VectorXi& perm, int nans, Eigen::VectorXi& ans)
{
    double rU;
    int i, j, n;
	n = p.size();
	// std::cout << "the value of n is " << n  << std::endl;
    int nm1 = n - 1;
	std::random_device rd;  // Will be used to obtain a seed for the random number engine
    std::mt19937 generator(rd()); // Standard mersenne_twister_engine seeded with rd()	
	std::uniform_real_distribution<double> distribution(0.0,1.0);

    /* record element identities */
    for (i = 0; i < n; i++)
        perm(i) = i + 1;

    /* sort the probabilities into descending order */
    // revsort(p, perm, n);
	sort_indexes(p, perm);
	std::sort(p.begin(),p.end());

    /* compute cumulative probabilities */
    for (i = 1 ; i < n; i++)
        p(i) += p(i - 1);

    for (i = 0; i < nans; i++) {
        rU = distribution(generator);		
        for (j = 0; j < nm1; j++) {
            if (rU <= p(j))
                break;
            }
        ans(i) = perm(j);
    }
}

void gom::generateGemanMcClureApprox1(Eigen::VectorXd& rjj,double sig,Eigen::MatrixXd& z1,Eigen::MatrixXd& z2)
{
	z1.setZero();
	z2.setZero();
	int N1 = rjj.size();
	Eigen::VectorXd sig1(N1),f11(N1),f12(N1),f21(N1),f22(N1), f111(N1),f211(N1);
	f11.setZero();
	f12.setZero();
	f21.setZero();
	f22.setZero();

	sig1 = sig*(rjj.array().inverse());
	sig1 = (rjj.array().isNaN()).select(std::numeric_limits<double>::infinity(),sig1); 

	Array3d sig_m;
	Array5d sig_l;

	sig_m << 0.6, 1.1, 2.1;
	sig_l <<0.3000, 0.6000, 1.6000, 2.3000, 3.7000;
	
	Array4d m_m, c_m;
	Array6d m_l, c_l; 
	
	m_m << -0.0000,   -0.1177,   -0.6082,   -1.4570;
	c_m << 0.0701,    0.3971,    0.8024,    0.9949;
	m_l << -0.0000,   -0.0054,   -0.0393,    0.1627,    0.4598,    0.7838;
	c_l << 0.0063,    0.0666,    0.1606,    0.0817,    0.0255,    0.0019;


	f11 = (sig1.array()>=sig_m(0)).select(m_m(1),f11);
	f11 = (sig1.array()>=sig_m(1)).select(m_m(2),f11);
	f11 = (sig1.array()>=sig_m(2)).select(m_m(3),f11);
	f111 =  f11.array()*(rjj.array().square());

	f12 = (sig1.array()>=sig_m(0)).select(c_m(1),f12);
	f12 = (sig1.array()>=sig_m(1)).select(c_m(2),f12);
	f12 = (sig1.array()>=sig_m(2)).select(c_m(3),f12);

	f21 = (sig1.array()>=sig_l(0)).select(m_l(1),f21);
	f21 = (sig1.array()>=sig_l(1)).select(m_l(2),f21);
	f21 = (sig1.array()>=sig_l(2)).select(m_l(3),f21);
	f21 = (sig1.array()>=sig_l(3)).select(m_l(4),f21);
	f21 = (sig1.array()>=sig_l(4)).select(m_l(5),f21);
	f211 = f21.array()*(rjj.array().square());

	f22 = (sig1.array()>=sig_l(0)).select(c_l(1),f22);
	f22 = (sig1.array()>=sig_l(1)).select(c_l(2),f22);
	f22 = (sig1.array()>=sig_l(2)).select(c_l(3),f22);
	f22 = (sig1.array()>=sig_l(3)).select(c_l(4),f22);
	f22 = (sig1.array()>=sig_l(4)).select(c_l(5),f22);
	
	z1 << f111, f12;
	z2 << f211, f22;
}

double gom::getMinEigVal(const Eigen::MatrixXd& H)
{
	Eigen::EigenSolver<MatrixXd> es(H);
	return (double)(es.eigenvalues().matrix().real().minCoeff());
}

Eigen::Matrix3d gom::vec2crossmat(Eigen::Vector3d v) {
    Eigen::Matrix3d C(3, 3);
    C.setZero();
    C(0, 1) = -v[2];
    C(0, 2) = v[1];
    C(1, 0) = v[2];
    C(1, 2) = -v[0];
    C(2, 0) = -v[1];
    C(2, 1) = v[0];
    return C;
}

void gom::w2R(Eigen::Vector3d& w, Eigen::Matrix3d& R) {
    double omega = w.norm();
    R.setIdentity();
    if(omega) {
        Eigen::Vector3d n = w/omega;
        double s = sin(omega);
        double c = cos(omega);
        double cc = 1 - c;
        double n0 = n(0),           n1 = n(1),          n2 = n(2);
        double n01cc = n0*n1*cc,    n12cc = n1*n2*cc,   n20cc = n2*n0*cc;
        double n0s = n0*s,          n1s = n1*s,         n2s = n2*s;

        R(0,0) = c+n0*n0*cc;    R(0,1) = n01cc-n2s;     R(0,2) = n20cc+n1s;
        R(1,0) = n01cc+n2s;     R(1,1) = c+n1*n1*cc;    R(1,2) = n12cc-n0s;
        R(2,0) = n20cc-n1s;     R(2,1) = n12cc+n0s;     R(2,2) = c+n2*n2*cc;
    }
}

Eigen::Matrix3d gom::getRandomRotation()
{
	Eigen::Vector3d v;
	Eigen::Matrix3d R;
    std::random_device rd;
    std::mt19937 e2(rd());
    std::normal_distribution<double> dist(0, 1);
	v(0) = dist(e2);
	v(1) = dist(e2);
	v(2) = dist(e2);
	double t = v.norm();
	v /= t;
	v *= (double)150.0*M_PI/180.0;
	w2R(v,R);
	return R;
}

double CApp::RobustUmeyamaPairwise()
{
	// Eigen::initParallel();
	printf("Starting robust umeyama for pairwise 3D registration\n");
	/* Dimensions check, already done */
	
	auto t_start =  std::chrono::high_resolution_clock::now();
	auto t_stop = std::chrono::high_resolution_clock::now();
	double time_fgr;

	t_start = std::chrono::high_resolution_clock::now();
	int N = num_corr_;		
	// Copying points into local variables
	PointsMat a = ptsB_;
	PointsMat b = ptsA_;
	/* Initialization of variables */
	double sig = std::numeric_limits<double>::infinity();
	// Assuming both pointclouds have the same size.

	int Iter = 0;
	int breakpoint=0;
	VectorXb ind_bool(N);
	ind_bool.setConstant(true);
	double sig_prev = -100.0d;
	double sig1 = sig;
	double sig_prev_Iter;
	double init_sig_start = sig_min_;
	double init_stepsize = 1;

	bool first_time = true;
	bool all_eig_neg = false;

	double perturb_rotangle = std::numeric_limits<double>::quiet_NaN();
	double perturb_trans = std::numeric_limits<double>::quiet_NaN();

    Eigen::VectorXd weights = Eigen::VectorXd::Constant(N,1,1);
	Eigen::Matrix4d transf = WeightedUmeyama(b, a, weights); 
	Eigen::Matrix3d R = transf.block<3,3>(0,0);
	Eigen::Vector3d T = transf.block<3,1>(0,3);

	// Declaration of all variables which are used inside the loop.
	PointsMat p, r;  
	int n = std::min(100,N); 
	int n_iter = 0;
	int N1=n;
	double start, stop, stepsize, lstart, lstop, temp, lmid; 
	Eigen::VectorXi sample_inds = Eigen::VectorXi::LinSpaced(N,0,N);
	Eigen::VectorXi perm = Eigen::VectorXi::LinSpaced(N,1,N);
	Eigen::VectorXd weights1(N); //, weights2(N+1);
	Eigen::VectorXd rjj(N,1);
	Eigen::MatrixXd z11(N,2),z22(N,2), z1(N,2), z2(N,2);
	Eigen::Matrix<double,6,6> A, B, grad_iit,H, A1, B1;
	Eigen::Matrix<double,6,1> grad_i;
	std::vector<Eigen::Matrix<double,6,1>> grad_i_all(N);
	Eigen::VectorXd sig2, l1;
	Eigen::Matrix3d Rprev,R2, rrT, H11, H12, H21, H22 = 2*Eigen::MatrixXd::Identity(3,3);
	Eigen::Vector3d Tprev,T2;

	PointsMat a1(3,n), b1(3,n), p1(3,n), r1(3,n);
	Eigen::Tensor<double,2> Rt(3,3);
	
	/* While loop start */
	#pragma omp declare reduction (+: Eigen::Matrix<double,6,6>: omp_out=omp_out+omp_in)\
		initializer(omp_priv=Eigen::Matrix<double,6,6>::Zero(6,6))	
	while(sig > sig_min_ && Iter<max_iters_)
	{
		sig_prev_Iter = sig;
		Iter += 1;
		p = a.colwise() - T;
		r = p - R*b;

		
		Rt = Eigen::TensorMap<Eigen::Tensor<double,2>>(R.data(),3,3); 
		if(stoch_sampling_)
		{
			if(Iter>=5)
			{
				weights1 = weights/weights.sum();
				sample_inds.resize(n);
				ProbSampleReplace(weights1, perm, n, sample_inds);
			}
			else
				sample_inds = Eigen::VectorXi::LinSpaced(N,0,N);
		}
		else
			sample_inds = Eigen::VectorXi::LinSpaced(N,0,N);
		
		sig_prev = sig;
		if(Iter==1)
			sig_prev = sig_max_;
		if(!fixed_anneal_)
		{
			N1 = sample_inds.size();
			a1 = a(Eigen::placeholders::all, sample_inds);
			b1 = b(Eigen::placeholders::all, sample_inds);
			p1 = p(Eigen::placeholders::all, sample_inds);
			r1 = r(Eigen::placeholders::all, sample_inds);

			rjj = r1.colwise().norm(); 

			z1.resize(N1,2);
			z2.resize(N1,2);
			generateGemanMcClureApprox1(rjj,sig_prev,z1,z2);
			A.setConstant(0); 
			B.setConstant(0);

			// #pragma omp parallel for if(N1>512) reduction(+:A,B) private(H11,H12,H,grad_i,grad_iit) //num_threads(32)// Assuming user-defined reductions aren't allowed
			for(int j=0; j < N1; j++)
			{
				H11 = 2*((p1(Eigen::placeholders::all, j).transpose()).dot(R*b1(Eigen::placeholders::all,j)))*Eigen::MatrixXd::Identity(3,3)-b1(Eigen::placeholders::all,j)*(p1(Eigen::placeholders::all,j).transpose())*R-R.transpose()*p1(Eigen::placeholders::all,j)*(b1(Eigen::placeholders::all,j).transpose());
				H12 = 2*vec2crossmat(b1(Eigen::placeholders::all,j))*R.transpose();
				
				H << H11, H12,
					H12.transpose(), H22;				
				grad_i << -H12*r1(Eigen::placeholders::all,j),
							-2*r1(Eigen::placeholders::all,j);
				grad_iit = (grad_i*grad_i.transpose())/(rjj(j)*rjj(j)); 
				A += z1(j,0)*H-4*z2(j,0)*grad_iit;
				B += z1(j,1)*H-4*z2(j,1)*grad_iit;
			}

			/* Search for sigma*/
			start = init_sig_start;
			stop=std::min(100.0d, sig_prev);
			stepsize=std::max(sig_min_/10,init_stepsize/pow(2,(Iter-1)));

			lstart = getMinEigVal(A/(start*start) + B);
			lstop = getMinEigVal(A/(stop*stop)+B);
			n_iter = 0;
			while((stop-start)>stepsize)
			{
				temp = (start+stop)/2.0d;
				lmid = getMinEigVal(A/(temp*temp)+B);
                if(lstart<eig_min_thres_ && lmid>eig_min_thres_) 
                {
					stop = temp;
                    lstop = lmid;
				}
                else if(lstart<eig_min_thres_ && lmid<=eig_min_thres_ && lstop>eig_min_thres_)
             	{
					start = temp;
                    lstart = lmid;
				}					
                else if(lstart>=eig_min_thres_)
                {
					stop = start;
                    lstart = lmid;
                    lstop = lmid;
				}
                else if(lstart<eig_min_thres_ && lmid<eig_min_thres_ && lstop<eig_min_thres_)
				{
                    stop = 100;
                    start = sig_min_; //% Could be modified to 0.1 %FIRST CHANGE
                    lstart = getMinEigVal((A/(start*start)+B)); lstop = getMinEigVal((A/(stop*stop)+B));
                    if(lstop<eig_min_thres_)
					{
						sig2 = Eigen::VectorXd::LinSpaced(100,0.1,10);
                        l1 = Eigen::VectorXd::Zero(100);
                        for(int k=0; k<100; k++)
                        {
							l1(k) = getMinEigVal((A/(sig2(k)*sig2(k))+B));
                            if(l1(k)>eig_min_thres_)
                            {
								stop = sig2(k);
                                start = sig2(k);
                                break;
							}
						}
                        if(l1.maxCoeff()<eig_min_thres_)
						{
                            stop = 0.9*sig_prev_Iter;
                            start = 0.9*sig_prev_Iter; 
                            n_times_inner_ += 1;
						}
                    }
                }
                n_iter = n_iter+1;
			}
			sig = stop;

            if((abs(sig_prev_Iter-sig)/sig_prev_Iter)<0.01 && sig>1)
			{
                sig = 0.9*sig_prev_Iter;
                n_times_inner_ += 1;
			}
            if(((sig>sig_prev_Iter && abs(sig-sig_prev_Iter)/sig_prev_Iter>thres_)))
            {
			    sig = 0.9*sig_prev_Iter;
                n_times_inner_ += 1;
			}        
            if(abs(sig-sig_prev)<1e-4 && stepsize<=sig_min_/10) 
            {
				if(first_time)
				{
                    sig = 0.9*sig_prev_Iter;
					n_times_inner_ += 1;
                    first_time = false;
				}
				else
				    break;
			}
		}
		else
			sig = sig_prev/anneal_factor_;


		if(!isumeya_)
		{
			// DO NOTHING, IDEALLY HAVE TO RUN SE3REG.
		}
		else
		{
			Rprev.array() = R.array()+5.0d; Tprev.array() = T.array()+5.0d;
			int tempo = 0;
			while(((Rprev-R).norm()>1e-2 || (Tprev-T).norm()>1e-3)) 
			{
				p = a.colwise() - T;
				r = p - R*b;	
				rjj = r.colwise().norm();
				weights = inverse(Eigen::pow((1+Eigen::pow(rjj.array()/sig,2)),2));
				Rprev = R; Tprev = T;
				transf = WeightedUmeyama(b, a, weights);
				R = transf.block<3,3>(0,0);
				T = transf.block<3,1>(0,3);				
				n_runs_ += 1;
			}
		}
		if(isperturb_)
		{
			/* perturb by a large rotation and translation*/
			Eigen::Matrix3d rand_rot = getRandomRotation();
			R2 = R*rand_rot;
			T2 = T;
			if(!isumeya_)
			{
				//DO NOTHING, ideally should run SE3Reg
			}
			else
			{
				Rprev.array() = R2.array()+5.0d; Tprev.array() = T2.array()+5.0d;
				int tempo = 0;
				while(((Rprev-R2).norm()>1e-2 || (Tprev-T2).norm()>1e-3)) //&& tempo<20
				{
					p = a.colwise() - T2;
					r = p - R2*b;	
					rjj = r.colwise().norm();
					weights = inverse(Eigen::pow((1+Eigen::pow(rjj.array()/sig,2)),2));
					Rprev = R2; Tprev = T2;
					transf = WeightedUmeyama(b, a, weights);
					R2 = transf.block<3,3>(0,0);
					T2 = transf.block<3,1>(0,3);				
					n_runs_ += 1;
				}				
			}
			auto r_temp1 =  (a - R*b).colwise() - T;
			auto r_norm1 = r_temp1.colwise().norm();
            double f_val1 = ((r_norm1.array().square()/(r_norm1.array().square()/(sig*sig)+1)).sum())/r_temp1.size();

			auto r_temp2 =  (a - R2*b).colwise() - T2;
			auto r_norm2 = r_temp2.colwise().norm();
            double f_val2 = ((r_norm2.array().square()/(r_norm2.array().square()/(sig*sig)+1)).sum())/r_temp2.size();

			if(abs(f_val2-f_val1)>1e-4 && f_val2 < f_val1)
			{
				R = R2; T = T2;
				n_perturbs_ += 1;
			}
		}

		if(Iter>2 && (sig_prev_Iter-sig)/sig_prev_Iter>0.5)
		{
			breakpoint += 1;
			continue;
		}

		if(breakpoint>0 && sig<=sig_prev_Iter && sig<sig_min_)
			break;
	}

	// Write the transformation to the private variable.	
	Eigen::Matrix4d M = Eigen::Matrix4d::Identity();
	M.block<3,3>(0,0) = R;
	M.block<3,1>(0,3) = T;
	TransOutput_ = M; //M.cast<float>();
	t_stop = std::chrono::high_resolution_clock::now();
	time_fgr = std::chrono::duration_cast<std::chrono::nanoseconds>(t_stop - t_start).count();
	std::cout << "The time taken for Robust Umeyama is " << time_fgr*1e-9 << std::endl;
	t_gom_ = time_fgr*1e-9;
	sig_stop_ = sig;
	n_iter_ = Iter;

	return (double)init_stepsize;
}

Eigen::Matrix4d CApp::GetOutputTrans()
{
	return TransOutput_;
}

